# !/usr/bin/env python
# -*- coding: utf-8 -*-
# @Time    : 2025/9/24 23:15
# @Author  : Dell
# @File    : emantic_route.py
# @Software: PyCharm
# @Desc    :语义路由
from langchain_openai import ChatOpenAI
from langchain.utils.math import cosine_similarity
from langchain_core.output_parsers import StrOutputParser
from langchain_core.prompts import PromptTemplate
from langchain_core.runnables import RunnableLambda, RunnablePassthrough
from langchain_community.embeddings import DashScopeEmbeddings
llm = ChatOpenAI(model_name='qwen-plus', base_url="https://dashscope.aliyuncs.com/compatible-mode/v1", temperature=0.5,api_key="sk-7b4f290b1a3e42a1a9a1957fa44eff37")
# 两个 prompts 物理教授、数据教授
physics_template = """你是一位非常聪明的物理学教授。\
你擅长以简洁易懂的方式回答有关物理的问题。\
当你不知道问题的答案时，你会承认你不知道。

这里有一个问题：
{query}"""

math_template = """你是一位非常优秀的数学家。你非常擅长回答数学问题。\
你之所以如此优秀，是因为你能将难题分解成各个组成部分，\
回答各个组成部分的问题，然后将它们组合起来回答更广泛的问题。

这里有一个问题：
{query}"""
# 嵌入
embeddings=DashScopeEmbeddings(
    model="text-embedding-v1", dashscope_api_key="sk-7b4f290b1a3e42a1a9a1957fa44eff37"
)
prompt_templates = [physics_template, math_template]
prompt_embeddings = embeddings.embed_documents(prompt_templates)

# 将问题转至提示
def prompt_router(input):
    # 嵌入问题
    query_embedding = embeddings.embed_query(input["query"])
    # 计算相似度
    similarity = cosine_similarity([query_embedding], prompt_embeddings)[0]
    most_similar = prompt_templates[similarity.argmax()]
    # 选择的提示
    print("呼叫数学专家" if most_similar == math_template else "呼叫物理专家")
    return PromptTemplate.from_template(most_similar)


chain = (
    {"query": RunnablePassthrough()}
    | RunnableLambda(prompt_router)
    | llm
    | StrOutputParser()
)

print(chain.invoke("什么是黑洞？"))
